import random
import string
from io import BytesIO
from PIL import Image, ImageDraw, ImageFont


class ImageCaptcha:
    def __init__(self):
        # 初始化字体（使用默认字体或指定字体文件路径）
        try:
            self.font = ImageFont.truetype("arial.ttf", size=36)
        except:
            # 如果找不到字体，使用默认字体
            self.font = ImageFont.load_default()

    def generate_captcha(self, width=120, height=40, char_num=4):
        """
        生成图片验证码
        :param width: 图片宽度
        :param height: 图片高度
        :param char_num: 验证码字符数量
        :return: (验证码文本, 图片字节流)
        """
        # 生成随机字符（字母+数字）
        chars = string.digits
        captcha_text = ''.join(random.choices(chars, k=char_num))

        # 创建图片
        image = Image.new('RGB', (width, height), (255, 255, 255))
        draw = ImageDraw.Draw(image)

        # 绘制干扰线
        for _ in range(5):
            x1 = random.randint(0, width)
            y1 = random.randint(0, height)
            x2 = random.randint(0, width)
            y2 = random.randint(0, height)
            draw.line([(x1, y1), (x2, y2)], fill=self._random_color(), width=1)

        # 绘制验证码字符
        for i, char in enumerate(captcha_text):
            x = 10 + i * (width // char_num)
            y = random.randint(0, height - 30)
            draw.text((x, y), char, font=self.font, fill=self._random_color())

        # 添加噪点
        for _ in range(100):
            x = random.randint(0, width)
            y = random.randint(0, height)
            draw.point((x, y), fill=self._random_color())

        # 转换为字节流
        buffer = BytesIO()
        image.save(buffer, format='PNG')
        buffer.seek(0)

        return captcha_text, buffer

    def _random_color(self):
        """生成随机颜色"""
        return (random.randint(50, 150),
                random.randint(50, 150),
                random.randint(50, 150)
                )